import getopt
import sys
import traceback
import json

import pandas as pd
from kafka import KafkaConsumer

from air_web.config.config import config
from air_web.data_platform import sql_engine


class ApprovalInfo():
    def __init__(self, reset):
        self.reset = reset

    def create_consumer(self):
        kafka_config = {
            'bootstrap_servers': ','.join(config['KAFKA_HOST']),  # 例如：'localhost:9092'
            'group_id': 'approval'
        }
        if self.reset:
            kafka_config['auto_offset_reset'] = 'earliest'

        # Create a Kafka consumer instance
        consumer = KafkaConsumer(config['TOPIC_AIR_CONDITION_APPROVAL'],
                                 value_deserializer=lambda v: json.loads(v),
                                 **kafka_config)
        return consumer

    def get_on5_dict(self):
        sql = """select org_no, org_name from real_org_no where org_level = 1"""
        on5_df = sql_engine.query(sql)
        on5_dict = on5_df.set_index('org_no')['org_name'].to_dict()
        return on5_dict

    def main(self):
        on5_dict = self.get_on5_dict()

        consumer = self.create_consumer()
        for message in consumer:
            data_dict = message.value
            try:
                data = data_dict['data']
                on5 = int(data['executeCityRange'])
                if on5 not in on5_dict:  # 过滤掉非地市的
                    print(f"{on5}非地市")
                    continue

                data['city_name'] = on5_dict[on5]
                data_df = pd.DataFrame([data])
                data_df = data_df.rename(
                    columns={'planId': 'plan_id', 'executeCityRange': 'on5',
                             'executeStartTime': 'control_time', 'executeEndTime': 'end_time',
                             'downIndex': 'conrol_prop', 'planType': 'plan_type'})

                data_df['conrol_prop'] = data_df['conrol_prop'].astype("float64") / 100
                data_df['contin_days'] = (pd.to_datetime(data_df['end_time']) - pd.to_datetime(
                    data_df['control_time'])).dt.days + 1

                sql_engine.update_df_by_id(data_df, config['CITY_APPROVAL_INFO'])
                print(f"写入mysql数据:{data_dict}")
            except:
                traceback.print_exc()
                print('error', data_dict)
        # Close Kafka consumer and MySQL connection
        consumer.close()


if __name__ == "__main__":
    reset = False
    opts, args = getopt.getopt(sys.argv[1:], "R")
    for opt, val in opts:
        if opt == "-R":
            reset = True

    ApprovalInfo(reset).main()

